#!/usr/bin/env python
# coding: utf-8

# # 回调机制 Callback
# 
# 在深度学习训练过程中，为及时掌握网络模型的训练状态、实时观察网络模型各参数的变化情况和实现训练过程中用户自定义的一些操作，luojianet_ms提供了回调机制（Callback）来实现上述功能。
# 
# Callback回调机制一般用在网络模型训练过程`Model.train`中，luojianet_ms的`Model`会按照Callback列表`callbacks`顺序执行回调函数，用户可以通过设置不同的回调类来实现在训练过程中或者训练后执行的功能。
# 
# > 更多内置回调类的信息及使用方式请参考[API文档](http://58.48.42.237/luojiaNet/luojiaNetapi/ #luojianet_ms.Callback)。
# 
# ## Callback介绍和使用
# 
# 当聊到回调Callback的时候，大部分用户都会觉得很难理解，是不是需要堆栈或者特殊的调度方式，实际上我们简单的理解回调：
# 
# 假设函数A有一个参数，这个参数是个函数B，当函数A执行完以后执行函数B，那么这个过程就叫回调。
# 
# `Callback`是回调的意思，luojianet_ms中的回调函数实际上不是一个函数而是一个类，用户可以使用回调机制来**观察训练过程中网络内部的状态和相关信息，或在特定时期执行特定动作**。
# 
# 例如监控损失函数Loss、保存模型参数ckpt、动态调整参数lr、提前终止训练任务等。
# 
# 下面以基于MNIST数据集训练LeNet-5网络模型为例，介绍几种常用的luojianet_ms内置回调类。
# 
# 首先需要下载并处理MNIST数据，构建LeNet-5网络模型，示例代码如下：

# In[1]:

import luojianet_ms.nn as nn
from luojianet_ms.train import Model
from luojianet_ms.dataset import MnistDataset
from src.lenet import lenet
from luojianet_ms.dataset.vision.c_transforms import RandomHorizontalFlip, HWC2CHW
from luojianet_ms.dataset.transforms.c_transforms import TypeCast
from luojianet_ms.common import dtype as mstype
epochs = 5
batch_size = 32
momentum = 0.9
learning_rate = 1e-2

dataset_train=MnistDataset(dataset_dir='./MNIST_data',usage='train')
# 数据增强
trans = [
    RandomHorizontalFlip(prob=0.5),
    HWC2CHW(),
    TypeCast(mstype.float32)
]

convert = [
    TypeCast(mstype.int32)
]

dataset_train = dataset_train.map(operations=trans, input_columns=['image'])
dataset_train = dataset_train.map(operations=convert, input_columns=['label']).batch(batch_size, True)

network = lenet()
net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
net_opt = nn.Momentum(network.trainable_params(), learning_rate=learning_rate, momentum=momentum)
model = Model(network, loss_fn=net_loss, optimizer=net_opt, metrics={'acc'})


# 回调机制的使用方法，在`model.train`方法中传入`Callback`对象，它可以是一个`Callback`列表，示例代码如下，其中ModelCheckpoint是luojianet_ms提供的回调类。

# In[2]:


from luojianet_ms.train.callback import LossMonitor,ModelCheckpoint

# 定义回调类
ckpt_cb = ModelCheckpoint()
loss_cb = LossMonitor(1875)

model.train(5, dataset_train, callbacks=[ckpt_cb, loss_cb])


# ## 常用的内置回调函数
# 
# luojianet_ms提供`Callback`能力，支持用户在训练/推理的特定阶段，插入自定义的操作。
# 
# ### ModelCheckpoint
# 
# 为了保存训练后的网络模型和参数，方便进行再推理或再训练，luojianet_ms提供了ModelCheckpoint配合使用。
# 
# 下面我们通过一段示例代码来说明如何保存训练后的网络模型和参数：

# In[3]:


from luojianet_ms.train.callback import CheckpointConfig,ModelCheckpoint

# 设置保存模型的配置信息
config_ck = CheckpointConfig(save_checkpoint_steps=1875, keep_checkpoint_max=10)
# 实例化保存模型回调接口，定义保存路径和前缀名
ckpoint = ModelCheckpoint(prefix="lenet", directory="./lenet", config=config_ck)

# 开始训练，加载保存模型和参数回调函数
model.train(1, dataset_train, callbacks=[ckpoint,loss_cb], dataset_sink_mode=False)


# 上面代码运行后，生成的Checkpoint文件目录结构如下：
# 
# ```text
# ./lenet/
# ├── lenet-1_1875.ckpt # 保存参数文件
# └── lenet-graph.meta # 编译后的计算图
# ```
# 
# ### LossMonitor
# 
# 为了监控训练过程中的损失函数值Loss变化情况，观察训练过程中每个epoch、每个step的运行时间，[luojianet_ms Vision](https://luojianet_ms.cn/vision/docs/zh-CN/r0.1/index.html)提供了`LossMonitor`接口（与luojianet_ms提供的`LossMonitor`接口有区别）。
# 
# 下面我们通过示例代码说明：

# In[4]:


from luojianet_ms.train.callback import LossMonitor

# 开始训练，加载保存模型和参数回调函数，LossMonitor的入参0.01为学习率，375为步长
model.train(5, dataset_train, callbacks=[LossMonitor(375)])



# 从上面的打印结果可以看出，luojianet_ms 提供的`LossMonitor`接口打印信息更加详细。由于步长设置的是375，所以每375个step会打印一条，loss值会波动，但总体来说loss值会逐步减小，精度逐步提高。
# 
# ### ValAccMonitor
# 
# 为了在训练过程中保存精度最优的网络模型和参数，需要边训练边验证，luojianet_ms Vision提供了`ValAccMonitor`接口。
# 
# 下面我们通过一段示例来介绍：

# In[5]:


from  src.EvalCallBack import EvalCallBack

dataset_eval=MnistDataset(dataset_dir='./MNIST_data',usage='test')
dataset_eval = dataset_eval.map(operations=trans, input_columns=['image'])
dataset_eval= dataset_eval.map(operations=convert, input_columns=['label']).batch(batch_size, True)


# 训练模型
def apply_eval(eval_param):
    eval_model = eval_param['model']
    eval_ds = eval_param['dataset']
    metrics_name = eval_param['metrics_name']
    res = eval_model.eval(eval_ds)
    return res[metrics_name]
eval_param_dict = {"model":model,"dataset":dataset_eval,"metrics_name":"acc"}
eval_cb = EvalCallBack(apply_eval, eval_param_dict)
model.train(1, dataset_train,callbacks=[eval_cb], dataset_sink_mode=True)


# 上面代码执行后，精度最优的网络模型和参数会被保存在当前目录下，文件名为"best.ckpt"。
# 
# ## 自定义回调机制
# 
# luojianet_ms不仅有功能强大的内置回调函数，当用户有自己的特殊需求时，还可以基于`Callback`基类自定义回调类。
# 
# 用户可以基于`Callback`基类，根据自身的需求，实现自定义`Callback`。`Callback`基类定义如下所示：

# In[6]:


class Callback():
    """Callback base class"""
    def begin(self, run_context):
        """Called once before the network executing."""
        pass # pylint: disable=W0107

    def epoch_begin(self, run_context):
        """Called before each epoch beginning."""
        pass # pylint: disable=W0107

    def epoch_end(self, run_context):
        """Called after each epoch finished."""
        pass # pylint: disable=W0107

    def step_begin(self, run_context):
        """Called before each step beginning."""
        pass # pylint: disable=W0107

    def step_end(self, run_context):
        """Called after each step finished."""
        pass # pylint: disable=W0107

    def end(self, run_context):
        """Called once after network training."""
        pass # pylint: disable=W0107


# 回调机制可以把训练过程中的重要信息记录下来，通过把一个字典类型变量`RunContext.original_args()`，传递给Callback对象，使得用户可以在各个自定义的Callback中获取到相关属性，执行自定义操作，也可以自定义其他变量传递给`RunContext.original_args()`对象。
# 
# `RunContext.original_args()`中的常用属性有：
# 
# - epoch_num：训练的epoch的数量
# - batch_num：一个epoch中step的数量
# - cur_epoch_num：当前的epoch数
# - cur_step_num：当前的step数
# 
# - loss_fn：损失函数
# - optimizer：优化器
# - train_network：训练的网络
# - train_dataset：训练的数据集
# - net_outputs：网络的输出结果
# 
# - parallel_mode：并行模式
# - list_callback：所有的Callback函数
# 
# 通过下面两个场景，我们可以增加对自定义Callback回调机制功能的了解。
# 
# ### 自定义终止训练
# 
# 实现在规定时间内终止训练功能。用户可以设定时间阈值，当训练时间达到这个阈值后就终止训练过程。
# 
# 下面代码中，通过`run_context.original_args`方法可以获取到`cb_params`字典，字典里会包含前文描述的主要属性信息。
# 
# 同时可以对字典内的值进行修改和添加，在`begin`函数中定义一个`init_time`对象传递给`cb_params`字典。每个数据迭代结束`step_end`之后会进行判断，当训练时间大于设置的时间阈值时，会向run_context传递终止训练的信号，提前终止训练，并打印当前的epoch、step、loss的值。

# In[7]:
import time
from luojianet_ms.train.callback import Callback

class StopTimeMonitor(Callback):

    def __init__(self, run_time):
        """定义初始化过程"""
        super(StopTimeMonitor, self).__init__()
        self.run_time = run_time            # 定义执行时间

    def begin(self, run_context):
        """开始训练时的操作"""
        cb_params = run_context.original_args()
        cb_params.init_time = time.time()   # 获取当前时间戳作为开始训练时间
        print("Begin training, time is:", cb_params.init_time)

    def step_end(self, run_context):
        """每个step结束后执行的操作"""
        cb_params = run_context.original_args()
        epoch_num = cb_params.cur_epoch_num  # 获取epoch值
        step_num = cb_params.cur_step_num    # 获取step值
        loss = cb_params.net_outputs         # 获取损失值loss
        cur_time = time.time()               # 获取当前时间戳

        if (cur_time - cb_params.init_time) > self.run_time:
            print("End training, time:", cur_time, ",epoch:", epoch_num, ",step:", step_num, ",loss:", loss)
            run_context.request_stop()       # 停止训练

model.train(5, dataset_train, callbacks=[LossMonitor( 1875), StopTimeMonitor(15)], dataset_sink_mode=False)



# 从上面的打印结果可以看出，当第3个epoch的第4673个step执行完时，运行时间到达了阈值并结束了训练。
# 
# ### 自定义阈值保存模型
# 
# 该回调机制实现当loss小于设定的阈值时，保存网络模型权重ckpt文件。
# 
# 示例代码如下：

# In[8]:


import luojianet_ms as ms

# 定义保存ckpt文件的回调接口
class SaveCkptMonitor(ms.train.callback.Callback):
    """定义初始化过程"""

    def __init__(self, loss):
        super(SaveCkptMonitor, self).__init__()
        self.loss = loss  # 定义损失值阈值

    def step_end(self, run_context):
        """定义step结束时的执行操作"""
        cb_params = run_context.original_args()
        cur_loss = cb_params.net_outputs.asnumpy() # 获取当前损失值

        # 如果当前损失值小于设定的阈值就停止训练
        if cur_loss < self.loss:
            # 自定义保存文件名
            file_name = str(cb_params.cur_epoch_num) + "_" + str(cb_params.cur_step_num) + ".ckpt"
            # 保存网络模型
            ms.save_checkpoint(save_obj=cb_params.train_network, ckpt_file_name=file_name)
            print("Saved checkpoint, loss:{:8.7f}, current step num:{:4}.".format(cur_loss, cb_params.cur_step_num))

model.train(1, dataset_train, callbacks=[SaveCkptMonitor(0.1)],dataset_sink_mode=False)


# 保存目录结构如下：
# 
# ```text
# ./
# ├── 1_253.ckpt
# ├── 1_258.ckpt
# ├── 1_265.ckpt
# ├── 1_332.ckpt
# ├── 1_358.ckpt
# ├── 1_380.ckpt
# ├── 1_395.ckpt
# ├── 1_1151.ckpt
# ├── 1_1358.ckpt
# ├── 1_1524.ckpt
# ```
